import cv2
import numpy as np
from ultralytics import YOLO




def get_v8_detect_onepic(model_path = 'yolov8n.pt'):
    # model_path = 'yolov8n.pt'
    model = YOLO(model_path)
    # model.predictor(imgsz=320, conf=0.5)
    model.predict(np.zeros((640,640,3),dtype=np.uint8), imgsz=640, conf=0.5, iou=0.5, device='cuda:0') # 第一次速度慢，先推理一次
    def detect_onepic(img, realtime_conf_thres=0.5, iou_thres=0.45, roi=None,pad=0,show_img=None):
        '''
        return:
        show_img:
        out: [[xyxy, cls_name, conf],]
        '''
        H, W = img.shape[:2]
        if roi is not None:
            x1, y1, x2, y2 = roi
            roi = [int(max(x1 - pad, 0)), int(max(y1 - pad, 0)), int(min(x2 + pad, W)), int(min(y2 + pad, H))]
        img2 = img[roi[1]:roi[3],roi[0]:roi[2],...] if roi is not None else img

        results = model.predict(img2, imgsz=640, conf=realtime_conf_thres, iou=iou_thres, device='cuda:0')
        # results = model.predict(img, imgsz=640, conf=realtime_conf_thres, iou=iou_thres, device='cuda:0')
        # print(results)

        # Process results list
        result = results[0] # one img
        boxes = result.boxes  # Boxes object for bounding box outputs
        # masks = result.masks  # Masks object for segmentation masks outputs
        # keypoints = result.keypoints  # Keypoints object for pose outputs
        # probs = result.probs  # Probs object for classification outputs
        # result.show()  # display to screen
        # result.save(filename='result.jpg')  # save to disk
        boxes = boxes.cpu().numpy()
        out = []
        for i in range(boxes.shape[0]):
            xyxy = boxes.xyxy[i, :].tolist()
            conf = boxes.conf[i]
            cls = int(boxes.cls[i])
            cls_name = result.names[cls]
            if roi is not None:  # 缩放回原图
                xyxy[0] += roi[0]
                xyxy[1] += roi[1]
                xyxy[2] += roi[0]
                xyxy[3] += roi[1]
            out.append([xyxy, cls_name, conf])

        # draw
        if show_img is None: #
            show_img = img.copy()
        img_diag = np.sqrt(show_img.shape[0] ** 2 + show_img.shape[1] ** 2)
        for xyxy, cls_name, conf in out:
            px1, py1, px2, py2 = int(xyxy[0]), int(xyxy[1]), int(xyxy[2]), int(xyxy[3])
            cv2.putText(show_img, f'{str(np.round(conf,2))}_{cls_name}', (px1, py1), cv2.FONT_HERSHEY_COMPLEX_SMALL, max(2, int(img_diag * 0.0005)), (0, 255, 0), max(2, int(img_diag * 0.0005)))
            # cv2.rectangle(show_img)
            cv2.rectangle(show_img, (px1, py1), (px2, py2), (0, 255, 0), max(2, int(img_diag * 0.001)))

        return show_img, out

    return detect_onepic


def ttest_get_detect_onepic():
    img = cv2.imread(r"D:\data\231215安全带\trainV8Det_flball_blue\_add_imgs\20240326-TRIAL1 4LINES\img_20240326_122555_422_NOREAD_1_OK_img.jpg")
    model_path = r"D:\data\231215安全带\trainV8Det_flball_blue\models\train7\weights\best.pt"
    detect_onepic = get_v8_detect_onepic(model_path)
    img, out = detect_onepic(img,0.25, roi=(1290,1460,2110,2320))
    print(out)

def ttest_dir():
    from ultralytics import YOLO
    test_dir = r"D:\data\231215安全带\trainV8Det_flball_blue\format_data\images\train"
    # Load a pretrained YOLOv8n model
    model_path = r"D:\data\231215安全带\trainV8Det_flball_blue\models\train7\weights\best.pt"
    model = YOLO(model_path)
    # Run inference on 'bus.jpg' with arguments
    model.predict(test_dir, save=True, imgsz=640, conf=0.5, show=True)

if __name__ == '__main__':
    ttest_get_detect_onepic()
    # ttest_dir()
